import io
import fitz
from PIL import Image
import requests
from tqdm import tqdm
import os

UAT = os.environ.get("UAT")  # 从环境变量中获取UAT，避免将UAT写在代码中

def pillow_image_to_file_binary(image):
    btyes_io = io.BytesIO()
    image.save(btyes_io, format='PNG')
    return btyes_io.getvalue()


def convert_pdf_to_images(pdf_binary, dpi=100):
    doc = fitz.open("pdf", pdf_binary)
    images = []
    for i in range(doc.page_count):
        page = doc[i]
    image = page.get_pixmap(dpi=dpi)
    image = Image.frombytes("RGB", [image.width, image.height], image.samples)
    images.append(image)
    return images


def pdf_ocr(image):
    api_url = "https://server.simpletex.cn/api/doc_ocr/"
    header = {"token": UAT}  # 鉴权信息，此处使用UAT方式
    img_file = {"file": pillow_image_to_file_binary(image)}
    res = requests.post(api_url, files=img_file, data={}, headers=header).json()  # 使用requests库上传文件
    print(res)
    return res["res"]["content"]


if __name__ == '__main__':
    pdf_path = '新高考数学知识点汇总.pdf'  # 输入pdf文件

    file_binary = open(pdf_path, 'rb').read()
    images = convert_pdf_to_images(file_binary)
    final_markdown_content = ""
    for image in tqdm(images):
        final_markdown_content += pdf_ocr(image) + "\n"

    open("input.md", "w", encoding="utf-8").write(final_markdown_content)
    print(final_markdown_content)  # 保存并输出最终markdown文件